import sys
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit
# Import functions from our quantum detector module
from quantum_asteroid_detector import (
    load_model, load_image, reduce_dimension, QuantumKNNClassifier,
    create_quantum_circuit, quantum_similarity
)

def detect_asteroid_in_image(image_path, model_path="quantum_asteroid_model.pkl", n_qubits=2, verbose=True):
    """Detect if an asteroid is present in the given image using quantum computing."""
    # Load the trained model
    model = load_model(model_path)
    
    # Load and preprocess the image
    if verbose:
        print(f"Loading and preprocessing image: {image_path}")
    features = load_image(image_path)
    
    # Reduce dimensions for quantum processing
    if verbose:
        print(f"Reducing dimensions for {n_qubits}-qubit quantum processing...")
    reduced_features = reduce_dimension(features, n_qubits)
    
    # Create and visualize quantum circuit for this image
    if verbose:
        print("Creating quantum circuit for feature encoding...")
    quantum_circuit = create_quantum_circuit(reduced_features, n_qubits)
    
    # Save circuit diagram
    circuit_path = "quantum_circuit_detection.png"
    quantum_circuit.draw(output='mpl', filename=circuit_path)
    if verbose:
        print(f"Quantum circuit saved to {circuit_path}")
    
    # Make prediction using quantum similarity
    if verbose:
        print("Performing quantum similarity calculations...")
    prediction = model.predict([reduced_features])[0]
    
    # Calculate confidence based on similarity scores
    similarities = []
    for train_features in model.X_train:
        similarity = quantum_similarity(reduced_features, train_features, n_qubits)
        similarities.append(similarity)
    
    # Find k most similar samples and their labels
    k = model.k
    most_similar_indices = np.argsort(similarities)[-k:]
    most_similar_labels = [model.y_train[i] for i in most_similar_indices]
    
    # Calculate confidence as proportion of matching labels
    confidence = sum(1 for label in most_similar_labels if label == prediction) / k
    
    # Display result
    if verbose:
        print(f"\nQuantum Detection Results:")
        print(f"Image: {image_path}")
        print(f"Prediction: {'Asteroid detected' if prediction == 1 else 'No asteroid detected'}")
        print(f"Confidence: {confidence:.2f}")
        print(f"Number of qubits used: {n_qubits}")
    
    # Display the image with results
    img = Image.open(image_path)
    plt.figure(figsize=(10, 8))
    
    # Create a 2x1 subplot layout
    plt.subplot(2, 1, 1)
    plt.imshow(np.array(img), cmap='gray')
    plt.title(f"Prediction: {'Asteroid detected' if prediction == 1 else 'No asteroid detected'} (Confidence: {confidence:.2f})")
    plt.axis('off')
    
    # Show the quantum circuit in the second subplot
    plt.subplot(2, 1, 2)
    circuit_img = plt.imread(circuit_path)
    plt.imshow(circuit_img)
    plt.title("Quantum Circuit Used for Detection")
    plt.axis('off')
    
    plt.tight_layout()
    plt.savefig("asteroid_detection_result.png")
    if verbose:
        print(f"Detection visualization saved to asteroid_detection_result.png")
    plt.show()
    
    return prediction, confidence

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python detect_asteroid.py <image_path>")
        sys.exit(1)
    
    image_path = sys.argv[1]
    detect_asteroid_in_image(image_path)